# 导入必要的库和模块
import pickle
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from sklearn import datasets
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from tqdm import tqdm


def main() -> None:
    # 加载数字数据集（sklearn 自带 8x8 灰度手写数字集）
    digits = datasets.load_digits()
    X: np.ndarray = digits.data  # 形状: (n_samples, 64)
    y: np.ndarray = digits.target

    # 将数据集划分为训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(
        X,
        y,
        test_size=0.2,
        random_state=42,
        stratify=y,
    )

    # 初始化变量以存储最佳结果
    best_accuracy: float = -1.0
    best_k: int | None = None
    best_model: Pipeline | None = None

    # 存储每个 k 的准确率，用于可视化
    k_values = list(range(1, 41))
    accuracies: list[float] = []

    # 尝试从 1 到 40 的 k 值
    for k in tqdm(k_values, total=len(k_values)):
        # 使用标准化 + KNN 的流水线，保证距离度量表现更稳定
        model = Pipeline(
            steps=[
                ("scaler", StandardScaler()),
                ("knn", KNeighborsClassifier(n_neighbors=k)),
            ]
        )

        model.fit(X_train, y_train)
        y_pred = model.predict(X_test)
        acc = accuracy_score(y_test, y_pred)
        accuracies.append(acc)

        if acc > best_accuracy:
            best_accuracy = acc
            best_k = k
            best_model = model

    # 将最佳 KNN 模型保存到二进制文件
    assert best_model is not None and best_k is not None
    model_path = Path(__file__).with_name("best_knn_model.pkl")
    with model_path.open("wb") as f:
        pickle.dump(best_model, f)

    # 训练过程可视化并保存为PDF
    pdf_path = Path(__file__).with_name("accuracy_plot.pdf")
    plt.figure(figsize=(10, 6))
    
    # 绘制准确率折线图
    plt.plot(k_values, accuracies, 'b-', linewidth=2, label='Accuracy')
    
    # 添加红色垂直线标记最优K值
    plt.axvline(x=best_k, color='r', linestyle='-', linewidth=2)
    
    # 标记最优点
    plt.plot(best_k, best_accuracy, 'ro', markersize=8)
    
    # 添加文字标注
    plt.text(best_k, best_accuracy, f'k={best_k}, Accuracy={best_accuracy:.3f}', 
             fontsize=10, ha='left', va='bottom', color='red', weight='bold')
    
    # 设置坐标轴标签和标题
    plt.xlabel("k value", fontsize=12)
    plt.ylabel("Accuracy", fontsize=12)
    plt.title("Accuracy of different k values", fontsize=14, weight='bold')
    
    # 设置坐标轴范围
    plt.xlim(0, 40)
    plt.ylim(min(accuracies) - 0.005, max(accuracies) + 0.005)
    
    # 添加网格
    plt.grid(True, linestyle=':', alpha=0.6)
    
    # 调整布局并保存为PDF
    plt.tight_layout()
    plt.savefig(pdf_path, format='pdf', dpi=300, bbox_inches='tight')
    plt.close()

    # 打印最佳结果
    print(f"最佳准确率为 {best_accuracy:.3f}, 对应的K值为 {best_k}")
    print(f"模型已保存至: {model_path}")
    print(f"准确率图表已保存至: {pdf_path}")


if __name__ == "__main__":
    main()
